import os
import cv2
import numpy as np
import random
from tqdm import tqdm
import re
"""
用于合成数据集
能够生成元件尺度的旋转bbox，而不是引脚尺度的.
todo
"""
def load_image(file_path):
    """Load a single image from the given path."""
    return cv2.imread(file_path, cv2.IMREAD_UNCHANGED)

def rotate_and_resize(image, mask, depth, vis,angle, scale):
    """Rotate and resize an image, its mask, and its depth map."""
    def get_bbox_from_mask(mask):
        """
        从掩码图像中获取掩码部分的边界框。

        :param mask: 输入的掩码图像，形状为 (h, w)
        :return: 边界框信息 (x, y, width, height) xywh
        """
        # 找到掩码中非零像素的坐标
        rows, cols = np.nonzero(mask)
        # bw,bh=bgwh
        # sx,sy=shift

        if len(rows) == 0 or len(cols) == 0:
            return (0, 0, 0, 0)

        # 计算边界框的左上角坐标和宽高
        x = np.min(cols)
        y = np.min(rows)
        width = np.max(cols) - x + 1
        height = np.max(rows) - y + 1

        return (x+width/2, y+height/2, width, height)
    h, w = image.shape[:2]
    center = (w // 2, h // 2)
    ori_cp_bbox=get_bbox_from_mask(mask)
    M = cv2.getRotationMatrix2D(center, angle, scale)

    # Compute the bounding box for the rotated image
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])
    new_w = int((h * sin) + (w * cos))
    new_h = int((h * cos) + (w * sin))

    # Adjust the rotation matrix to account for the new dimensions
    M[0, 2] += (new_w / 2) - center[0]
    M[1, 2] += (new_h / 2) - center[1]

    # Perform the rotation and resizing
    rotated_image = cv2.warpAffine(image, M, (new_w, new_h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0))
    rotated_mask = cv2.warpAffine(mask, M, (new_w, new_h), flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_CONSTANT, borderValue=0)
    rotated_depth = cv2.warpAffine(depth, M, (new_w, new_h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=0)
    rotated_vis = cv2.warpAffine(vis, M, (new_w, new_h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0))
    obb=ori_cp_bbox*scale+(angle,)
    return rotated_image, rotated_mask, rotated_depth,rotated_vis

def is_non_overlapping(mask, component_mask, x, y):
    """Check if a component can be placed at (x, y) without overlapping."""
    ch, cw = component_mask.shape[:2]
    if x + cw > mask.shape[1] or y + ch > mask.shape[0]:
        return False

    region = mask[y:y + ch, x:x + cw]
    region_non_zero = region > 0
    component_mask_non_zero = component_mask > 0
    overlap = np.sum(np.logical_and(region_non_zero, component_mask_non_zero))
    # overlap = np.sum(cv2.bitwise_and(region > 0, component_mask > 0))
    return overlap == 0

def place_component(board, board_depth, board_mask, board_vis_mask, 
                    component, component_depth, component_mask, component_vis_mask,
                    x, y):
    """Place a component on the board and update depth, mask, and visualized mask."""
    ch, cw = component.shape[:2]
    foreground_indices = component_mask > 0

    for c in range(3):
        board[y:y + ch, x:x + cw, c][foreground_indices] = component[:, :, c][foreground_indices]

    board_depth[y:y + ch, x:x + cw][foreground_indices] = component_depth[foreground_indices]

    board_mask[y:y + ch, x:x + cw][foreground_indices] = np.maximum(
        board_mask[y:y + ch, x:x + cw][foreground_indices], component_mask[foreground_indices]
    )

    board_vis_mask[y:y + ch, x:x + cw][foreground_indices] = component_vis_mask[foreground_indices]
    # board_vis_mask[y:y + ch, x:x + cw][foreground_indices] = cv2.applyColorMap(
    #     component_mask[foreground_indices].astype(np.uint8) * (255 // 178), cv2.COLORMAP_JET
    # )[foreground_indices]


def synthesize_pcb_dataset(empty_board_path, empty_board_depth_path, component_root,component_depth_root,output_folder, num_images, min_scale=0.5, max_scale=1.5):
    """Generate synthetic PCB dataset."""
   
    def get_component_type(component_path):
        type_name=component_path[len(component_root):]
        type_name=os.path.dirname(type_name)
        pattern = r'JRS(\d+)[a-zA-Z].*'
        match = re.match(pattern, type_name)
        number = match.group(1)
        return int(number)
    
    # Load the empty board and depth
    empty_board = load_image(empty_board_path)
    empty_board_depth = load_image(empty_board_depth_path)
    bh,bw=empty_board_depth.shape
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    component_paths = []#len=1948
    for root, dirs, files in os.walk(component_root):
        for file in files:
            if file.endswith(('.png', '.jpg', '.jpeg')) and 'label' not in root:
                component_paths.append(os.path.join(root, file))
    
    max_n_cp=int(bh*bw/300/400*1.2)
    min_n_cp=int(bh*bw/300/400*0.5)
    for i in tqdm(range(num_images)):
        board = empty_board.copy()
        board_depth = empty_board_depth.copy()
        board_mask = np.zeros(board.shape[:2], dtype=np.uint8)
        vis_mask = np.zeros_like(board)
        yolo_labels=[]

        used_components = []
        num_components = random.randint(min_n_cp, max_n_cp)  # Random number of components per board
        error_guest=[
            "JRS183R0603/_a_SaveData__20241013161630397_Save Data [14837, 35647]",
            "JRS156R0603/_a_SaveData__20241013161630396_Save Data [22823, 35559]",
            "JRS157R0603/_a_SaveData__20241013161630397_Save Data [18766, 35613]",
            "JRS116PGLQFP/_a_SaveData__20241013155754673_SaveData [23447, 102811]",
            "JRS104TSOP50",
            "JRS116PGLQFP",
            "JRS82R1812/_a_SaveData__20241013155402990_Save Data [109592, 18847]",
        ]
        for it in range(num_components):
            component_path = random.choice(component_paths)
            label_path = component_path.replace('/_a_SaveData', '/label/_a_SaveData').replace('jpg','png')
            depth_path = component_path.replace(component_root, component_depth_root+'/').replace('_TrueLight.jpg','.png')
            vis_path = component_path.replace('/_a_SaveData', '/label/_a_SaveData').replace('.jpg','_pseudo.png')
            
            log_file_path = 'og.log'
            error_flag=False
            for error_path in error_guest:
                if error_path in component_path:
                    it-=1
                    error_flag=True
                    break
            if error_flag:
                continue
            
            
            component = load_image(component_path)
            component_mask = load_image(label_path)
            component_depth = load_image(depth_path)
            component_vis_mask=load_image(vis_path)


            # if component is None or component_mask is None or component_depth is None:
            #     print(f"Warning: Failed to load data for {component_path}")
            #     print(component is None , component_mask is None , component_depth is None)
            #     continue
            
            if component is None or component_mask is None or component_depth is None:
                warning_message = f"Warning: Failed to load data for {component_path}"
                status_message = f"{component is None}, {component_mask is None}, {component_depth is None}"


                with open(log_file_path, 'a') as log_file:
                    log_file.write(warning_message + '\n')
                    log_file.write(status_message + '\n')
                continue

            angle = random.uniform(0, 360)
            scale = random.uniform(min_scale, max_scale)

            component, component_mask, component_depth ,component_vis_mask= rotate_and_resize(
                component, component_mask, component_depth,component_vis_mask, angle, scale/2
            )

            ch, cw = component.shape[:2]
            do_it_more=max(1,1024*1024/ch/cw-1)
            for _ in range(100):  # Maximum attempts to place the component
                x = random.randint(0, board.shape[1] - cw)
                y = random.randint(0, board.shape[0] - ch)

                if is_non_overlapping(board_mask, component_mask, x, y):
                    bbox_xywh=get_bbox_from_mask(component_mask,(bw,bh),(x,y))
                    class_id=get_component_type(component_path)
                    place_component(
                        board, board_depth, board_mask, vis_mask,
                        component, component_depth, component_mask,component_vis_mask,
                        x, y
                    )
                    used_components.append(component_path)
                    yolo_labels.append((class_id,)+bbox_xywh)
                    break
                    
                    # if do_it_more>0:
                    #     do_it_more-=1
                    # else:
                    #     break

            del component, component_mask, component_depth,component_vis_mask

        # Save outputs
        dir_list=["image","depth","mask","vis","cplist","cpbbox"]
        for dir_name in dir_list:
            if not os.path.exists(os.path.join(output_folder,dir_name)):
                os.makedirs(os.path.join(output_folder,dir_name))
        output_image_path = os.path.join(os.path.join(output_folder,"image"), f"pcb_{i:04d}.png")
        output_depth_path = os.path.join(os.path.join(output_folder,"depth"), f"pcb_{i:04d}_depth.png")
        output_mask_path = os.path.join(os.path.join(output_folder,"mask"), f"pcb_{i:04d}_mask.png")
        output_vis_mask_path = os.path.join(os.path.join(output_folder,"vis"), f"pcb_{i:04d}_vis_mask.png")
        output_components_path = os.path.join(os.path.join(output_folder,"cplist"), f"pcb_{i:04d}_components.txt")
        output_bbox_path=os.path.join(os.path.join(output_folder,"cpbbox"), f"pcb_{i:04d}_cpbbox.txt")
        cv2.imwrite(output_image_path, board)
        cv2.imwrite(output_depth_path, board_depth)
        cv2.imwrite(output_mask_path, board_mask)
        cv2.imwrite(output_vis_mask_path, vis_mask)
        with open(output_bbox_path, 'w') as f2:
            for label in yolo_labels:
                bbox_label=label[:5]
                label_str = " ".join(str(i) for i in bbox_label)
                f2.write(label_str + "\n")

        with open(output_components_path, 'w') as f:
            f.write('\n'.join(used_components))

# Example usage
# empty_board_path="/home/lhx/pcbsyn/smallpcbb.jpg"
empty_board_path="pcbbgimg/pcbbg/pcbbgrgb5000.jpg"
# empty_board_depth_path="/home/lhx/pcbsyn/smallpcbbd.jpg"
empty_board_depth_path="pcbbgimg/pcbbgd/pcbbgd5000.jpg"
component_root="/home/lhx/pcbsyn/data_zoo/PCBA/20241209_PCBA样板标图数据-华科用/"
output_folder="synbboxpcb"
component_depth_root="/home/lhx/pcbsyn/pcbad"
synthesize_pcb_dataset(
    empty_board_path=empty_board_path,
    empty_board_depth_path=empty_board_depth_path,
    component_root=component_root,
    component_depth_root=component_depth_root,
    output_folder=output_folder,
    num_images=100,
    min_scale=1.1,
    max_scale=0.9
)
